import logging
import time
from typing import Dict, Any, List, Tuple
import os
import json
import traceback
import fcntl
import re
import functools

from web3 import Web3
from eth_account import Account
from eth_account.signers.local import LocalAccount

from coin.exchange.uniswap_base.constants import (
    TOKEN_TO_ADDRESS,
    TOKEN_TO_DECIMAL_MULTIPLIER,
)

TX_ONCHAIN_LAG_TIME = {
    "ethereum": 15,
    "arbitrum": 1,
}
NONCE_NEED_INITIALIZATION = 0
PRIVATE_KEY_FILE_PATH = os.path.expanduser("~/private_keys.json")
NONCE_FILE_PATH_PREFIX = os.path.expanduser("~/nonce_")

account_list = []
with open(PRIVATE_KEY_FILE_PATH) as fi:
    json_list = json.load(fi)
    for wallet in json_list:
        public_address = wallet['public_address']
        private_key = wallet['private_key']
        account_list.append(Account.from_key(private_key))

account_map = {account.address: account for account in account_list}


def read_nonce_from_file(nonce_file) -> int:
    if nonce_file is None:
        return NONCE_NEED_INITIALIZATION
    nonce_val = NONCE_NEED_INITIALIZATION
    try:
        nonce_file.seek(0)
        text = nonce_file.readline()
        if text:
            nonce_val = int(text)
    except Exception:
        logging.exception('Exception raised from read_nonce_from_file')
        nonce_val = NONCE_NEED_INITIALIZATION
    return nonce_val


def write_nonce_to_file(nonce_file, nonce_val: int):
    if nonce_file is None:
        return
    try:
        nonce_file.seek(0)
        nonce_file.truncate()
        nonce_file.write(str(nonce_val))
    except Exception:
        logging.exception('Exception raised from write_nonce_to_file')


def get_nonce(web3: Web3, public_address: str, nonce_file) -> int:
    """_summary_
    Args:
        web3 (_type_): perform RPC method to get transaction count of remote address, may change to different block_identifier settings
        public_address (_type_): public_address of the wallet

    Raises:
        NotImplementedError: the public_address of the wallet is not in the wallet pool this instance manages

    Returns:
        int: which is nonce type
    """
    logging.info("[Native client]get_nonce started")
    public_address = Web3.to_checksum_address(public_address)
    if public_address not in account_map.keys():
        raise ValueError(public_address)
    nonce_val = read_nonce_from_file(nonce_file)
    if nonce_val == NONCE_NEED_INITIALIZATION:
        logging.info("[Native client]query for nonce")
        nonce_val = web3.eth.get_transaction_count(
            public_address, block_identifier="pending"
        )
    logging.info("[Native client]get_nonce ended")
    return nonce_val


def get_account(public_address: str) -> LocalAccount:
    public_address = Web3.to_checksum_address(public_address) 
    if public_address in account_map.keys():
        return account_map.get(public_address, None)
    raise ValueError(public_address)


def fill_remaining_tx(web3: Web3, tx_param: Dict[str, Any], nonce_file) -> Dict[str, Any]:
    from_addr = tx_param['from']
    assert Web3.is_checksum_address(from_addr)

    tx_param |= {
        "nonce": get_nonce(web3, from_addr, nonce_file),
        "maxFeePerGas": 1200000000,
        "maxPriorityFeePerGas": 120000000
    }

    logging.info('after sign: %s' % tx_param)
    return tx_param


@functools.lru_cache()
def parse_symbol(symbol: str):
    result = re.match("^\s*(\w+)-(\w+)(.(\d+)bp)?\s*$", symbol)

    mresult = result.groups()

    match mresult:
        case 'ETH', quote_currency, _, '1' | '5' | '30' | '100':
            return 'WETH', quote_currency, int(mresult[3]) * 100

        case base_currency, 'ETH', _, '1' | '5' | '30' | '100':
            return base_currency, 'WETH', int(mresult[3]) * 100

        case base_currency, quote_currency, _, '1' | '5' | '30' | '100':
            return base_currency, quote_currency, int(mresult[3]) * 100

        case base_currency, quote_currency, *_:
            return base_currency, quote_currency
        case _:
            raise NotImplementedError()
    return None


def calculate_fill_price(receipt, symbol, blockchain_name) -> float or None:
    base_currency, quote_currency, fee = parse_symbol(symbol)
    base_qty = 0
    quote_qty = 0
    for log in receipt.logs:
        topics = log.topics
        if topics[0].hex() == '0xddf252ad1be2c89b69c2b068fc378daa952ba7f163c4a11628f55a4df523b3ef':
            token_address = log.address
            if TOKEN_TO_ADDRESS[blockchain_name][base_currency] == token_address:
                base_qty = int(log.data.hex(), 16) * TOKEN_TO_DECIMAL_MULTIPLIER[base_currency]
            elif TOKEN_TO_ADDRESS[blockchain_name][quote_currency] == token_address:
                quote_qty = int(log.data.hex(), 16) * TOKEN_TO_DECIMAL_MULTIPLIER[quote_currency]
    if base_qty > 0 and quote_qty > 0:
        return quote_qty / base_qty
    return None


def sign_and_wait(
    web3: Web3,
    tx_param_list: List[Dict[str, Any]],
    symbol: str or None,
    blockchain_name="ethereum",
) -> Tuple[bool, str or None, float or None]:
    # Currently only support single tx case.
    assert len(tx_param_list) == 1
    tx_param = tx_param_list[0]
    tx_param["from"] = Web3.to_checksum_address(tx_param["from"])
    assert "to" in tx_param.keys()
    account = get_account(tx_param["from"])
    
    nonce_file_path = NONCE_FILE_PATH_PREFIX + tx_param["from"]
    nonce_file = None
    try:
        nonce_file = open(nonce_file_path, 'a+')
        fcntl.flock(nonce_file, fcntl.LOCK_EX)
    except Exception:
        logging.exception('Exception raised from open nonce file.')
        if nonce_file is not None:
            fcntl.flock(nonce_file, fcntl.LOCK_UN)
            nonce_file.close()
        return False, None, None
    
    try:
        logging.info("[Native client]fill_remaining_tx started")
        fill_remaining_tx(web3, tx_param, nonce_file)
        logging.info("[Native client]sign_transaction started")
        signed_tx = account.sign_transaction(tx_param)
        try:
            logging.info("[Native client]send_raw_transaction started")
            tx_hash = web3.eth.send_raw_transaction(signed_tx.rawTransaction)
            logging.info("[Native client]wait_for_transaction_receipt started")
            web3.eth.wait_for_transaction_receipt(tx_hash, poll_latency=1)
        except Exception:
            traceback.print_exc()
            logging.info("[Native client Error]sign_and_wait: excpetion during send and wait tx")
            write_nonce_to_file(nonce_file, NONCE_NEED_INITIALIZATION)
            return False, None, None
        logging.info("[Native client]sleep 2 sec started")
        time.sleep(TX_ONCHAIN_LAG_TIME[blockchain_name] * 2)
        logging.info("[Native client]get_transaction_receipt started")
        rcpt = web3.eth.get_transaction_receipt(tx_hash)
        if rcpt["status"] == 0:
            logging.info("[Native client Error]sign_and_wait: rcpt status is 0!")
            write_nonce_to_file(nonce_file, NONCE_NEED_INITIALIZATION)
            return False, None, None
        else:
            write_nonce_to_file(nonce_file, int(tx_param['nonce']) + 1)
        logging.info("[Native client]sign_and_wait end")
        if symbol is None:
            return True, tx_hash.hex(), None
        else:
            return True, tx_hash.hex(), calculate_fill_price(rcpt, symbol, blockchain_name)
    finally:
        logging.info("[Native client]close nonce file")
        if nonce_file is not None:
            fcntl.flock(nonce_file, fcntl.LOCK_UN)
            nonce_file.close()
